#pragma once

#include "vkinfer/runtime/vk/vkop.h"

namespace vkinfer
{
    class VkOpConv : public VkOp
    {
    private:
        struct PushConstants
        {
            // input shape
            uint32_t input_c;
            uint32_t input_h;
            uint32_t input_w;
            uint32_t output_c;
            uint32_t output_h;
            uint32_t output_w;
            
            // op config
            // uint32_t dilation;
            // uint32_t group;
            uint32_t kernel;
            uint32_t pad;
            uint32_t stride;
        } pc;

        std::shared_ptr<VkBlob> weight;
        std::shared_ptr<VkBlob> bias;
        
    public:
        VkOpConv(VkDevice device, const std::shared_ptr<OnnxNodeInfo>& opinfo);
        virtual void set_weight(const std::string& name, const std::shared_ptr<Blob>& blob) override;
        virtual void update_descriptors(
            const std::vector<std::shared_ptr<VkBlob>>& inputs,
            std::vector<std::shared_ptr<VkBlob>>& outputs) override;
        virtual void forward(const std::vector<std::shared_ptr<Blob>>& inputs,
            std::vector<std::shared_ptr<Blob>>& outputs) override;
    };
}
